"""Message handling service for conversations."""

from typing import Optional, List, Dict, Any, AsyncGenerator
from datetime import datetime
import asyncio
import json
from bson import ObjectId

import motor.motor_asyncio
from fastapi import WebSocket, WebSocketDisconnect

from ..models.conversation_model import Message, MessageRole, Conversation


class MessageHandler:
    """Handles message operations and real-time updates."""
    
    def __init__(self, mongodb_client: motor.motor_asyncio.AsyncIOMotorClient):
        """Initialize message handler.
        
        Args:
            mongodb_client: MongoDB client instance
        """
        self.mongodb = mongodb_client
        self.db = mongodb_client.ai_writing
        self.conversations = self.db.conversations
        
        # WebSocket connection management
        self.active_connections: Dict[str, List[WebSocket]] = {}
    
    async def add_message(
        self,
        conversation_id: str,
        user_id: str,
        content: str,
        role: MessageRole = MessageRole.USER,
        model: Optional[str] = None,
        tokens_used: int = 0,
        cost: float = 0.0
    ) -> Optional[Message]:
        """Add a message to a conversation.
        
        Args:
            conversation_id: Conversation ID
            user_id: User ID for ownership verification
            content: Message content
            role: Message role (user/assistant/system)
            model: AI model used (for assistant messages)
            tokens_used: Tokens consumed
            cost: Cost of the message
            
        Returns:
            Created Message instance or None if failed
        """
        message = Message(
            role=role,
            content=content,
            model=model,
            tokens_used=tokens_used,
            cost=cost
        )
        
        # Update conversation with new message
        result = await self.conversations.update_one(
            {
                "_id": ObjectId(conversation_id),
                "user_id": user_id
            },
            {
                "$push": {"messages": message.to_dict()},
                "$inc": {
                    "statistics.total_messages": 1,
                    "statistics.total_tokens": tokens_used,
                    "statistics.total_cost": cost
                },
                "$set": {
                    "statistics.last_message_at": message.timestamp,
                    "updated_at": datetime.utcnow()
                }
            }
        )
        
        if result.modified_count > 0:
            # Broadcast to WebSocket connections
            await self.broadcast_message(conversation_id, message)
            return message
        
        return None
    
    async def get_messages(
        self,
        conversation_id: str,
        user_id: str,
        skip: int = 0,
        limit: int = 50
    ) -> List[Message]:
        """Get messages from a conversation.
        
        Args:
            conversation_id: Conversation ID
            user_id: User ID for ownership verification
            skip: Number of messages to skip
            limit: Maximum number of messages to return
            
        Returns:
            List of messages
        """
        # Use aggregation to get paginated messages
        pipeline = [
            {
                "$match": {
                    "_id": ObjectId(conversation_id),
                    "user_id": user_id
                }
            },
            {
                "$project": {
                    "messages": {
                        "$slice": ["$messages", skip, limit]
                    }
                }
            }
        ]
        
        cursor = self.conversations.aggregate(pipeline)
        result = await cursor.to_list(length=1)
        
        if result and result[0].get("messages"):
            return [Message.from_dict(msg) for msg in result[0]["messages"]]
        
        return []
    
    async def update_message(
        self,
        conversation_id: str,
        user_id: str,
        message_id: str,
        new_content: str
    ) -> bool:
        """Update a message content.
        
        Args:
            conversation_id: Conversation ID
            user_id: User ID for ownership verification
            message_id: Message ID to update
            new_content: New message content
            
        Returns:
            True if updated successfully
        """
        # First, get the current message to save edit history
        doc = await self.conversations.find_one(
            {
                "_id": ObjectId(conversation_id),
                "user_id": user_id,
                "messages.id": message_id
            },
            {"messages.$": 1}
        )
        
        if not doc or not doc.get("messages"):
            return False
        
        old_message = doc["messages"][0]
        
        # Update the message with edit history
        result = await self.conversations.update_one(
            {
                "_id": ObjectId(conversation_id),
                "user_id": user_id,
                "messages.id": message_id
            },
            {
                "$set": {
                    "messages.$.content": new_content,
                    "messages.$.metadata.edited": True,
                    "updated_at": datetime.utcnow()
                },
                "$push": {
                    "messages.$.metadata.edit_history": {
                        "content": old_message["content"],
                        "edited_at": datetime.utcnow()
                    }
                }
            }
        )
        
        if result.modified_count > 0:
            # Broadcast update
            updated_message = Message.from_dict(old_message)
            updated_message.content = new_content
            updated_message.metadata.edited = True
            await self.broadcast_message_update(conversation_id, updated_message)
            return True
        
        return False
    
    async def delete_message(
        self,
        conversation_id: str,
        user_id: str,
        message_id: str
    ) -> bool:
        """Delete a message from conversation.
        
        Args:
            conversation_id: Conversation ID
            user_id: User ID for ownership verification
            message_id: Message ID to delete
            
        Returns:
            True if deleted successfully
        """
        result = await self.conversations.update_one(
            {
                "_id": ObjectId(conversation_id),
                "user_id": user_id
            },
            {
                "$pull": {"messages": {"id": message_id}},
                "$inc": {"statistics.total_messages": -1},
                "$set": {"updated_at": datetime.utcnow()}
            }
        )
        
        if result.modified_count > 0:
            await self.broadcast_message_deletion(conversation_id, message_id)
            return True
        
        return False
    
    async def stream_message_update(
        self,
        conversation_id: str,
        user_id: str,
        message_id: str,
        content_stream: AsyncGenerator[str, None]
    ):
        """Stream message content updates for real-time generation.
        
        Args:
            conversation_id: Conversation ID
            user_id: User ID for ownership verification
            message_id: Message ID being updated
            content_stream: Async generator yielding content chunks
        """
        full_content = ""
        
        async for chunk in content_stream:
            full_content += chunk
            
            # Update message in database periodically
            if len(full_content) % 100 == 0:  # Update every 100 chars
                await self.conversations.update_one(
                    {
                        "_id": ObjectId(conversation_id),
                        "user_id": user_id,
                        "messages.id": message_id
                    },
                    {
                        "$set": {
                            "messages.$.content": full_content,
                            "updated_at": datetime.utcnow()
                        }
                    }
                )
            
            # Broadcast chunk to WebSocket connections
            await self.broadcast_message_chunk(conversation_id, message_id, chunk)
        
        # Final update with complete content
        await self.conversations.update_one(
            {
                "_id": ObjectId(conversation_id),
                "user_id": user_id,
                "messages.id": message_id
            },
            {
                "$set": {
                    "messages.$.content": full_content,
                    "updated_at": datetime.utcnow()
                }
            }
        )
    
    # WebSocket connection management
    
    async def connect_websocket(self, websocket: WebSocket, conversation_id: str):
        """Connect a WebSocket client to a conversation.
        
        Args:
            websocket: WebSocket connection
            conversation_id: Conversation ID to subscribe to
        """
        await websocket.accept()
        
        if conversation_id not in self.active_connections:
            self.active_connections[conversation_id] = []
        
        self.active_connections[conversation_id].append(websocket)
    
    async def disconnect_websocket(self, websocket: WebSocket, conversation_id: str):
        """Disconnect a WebSocket client from a conversation.
        
        Args:
            websocket: WebSocket connection
            conversation_id: Conversation ID to unsubscribe from
        """
        if conversation_id in self.active_connections:
            if websocket in self.active_connections[conversation_id]:
                self.active_connections[conversation_id].remove(websocket)
            
            # Clean up empty conversation rooms
            if not self.active_connections[conversation_id]:
                del self.active_connections[conversation_id]
    
    async def broadcast_message(self, conversation_id: str, message: Message):
        """Broadcast a new message to all connected clients.
        
        Args:
            conversation_id: Conversation ID
            message: Message to broadcast
        """
        if conversation_id in self.active_connections:
            message_data = {
                "type": "new_message",
                "message": message.to_dict()
            }
            
            # Send to all connected clients
            disconnected = []
            for websocket in self.active_connections[conversation_id]:
                try:
                    await websocket.send_json(message_data)
                except WebSocketDisconnect:
                    disconnected.append(websocket)
                except Exception:
                    disconnected.append(websocket)
            
            # Clean up disconnected clients
            for ws in disconnected:
                await self.disconnect_websocket(ws, conversation_id)
    
    async def broadcast_message_update(self, conversation_id: str, message: Message):
        """Broadcast a message update to all connected clients.
        
        Args:
            conversation_id: Conversation ID
            message: Updated message
        """
        if conversation_id in self.active_connections:
            message_data = {
                "type": "message_updated",
                "message": message.to_dict()
            }
            
            disconnected = []
            for websocket in self.active_connections[conversation_id]:
                try:
                    await websocket.send_json(message_data)
                except:
                    disconnected.append(websocket)
            
            for ws in disconnected:
                await self.disconnect_websocket(ws, conversation_id)
    
    async def broadcast_message_chunk(self, conversation_id: str, message_id: str, chunk: str):
        """Broadcast a message chunk for streaming updates.
        
        Args:
            conversation_id: Conversation ID
            message_id: Message ID being streamed
            chunk: Content chunk
        """
        if conversation_id in self.active_connections:
            chunk_data = {
                "type": "message_chunk",
                "message_id": message_id,
                "chunk": chunk
            }
            
            disconnected = []
            for websocket in self.active_connections[conversation_id]:
                try:
                    await websocket.send_json(chunk_data)
                except:
                    disconnected.append(websocket)
            
            for ws in disconnected:
                await self.disconnect_websocket(ws, conversation_id)
    
    async def broadcast_message_deletion(self, conversation_id: str, message_id: str):
        """Broadcast a message deletion to all connected clients.
        
        Args:
            conversation_id: Conversation ID
            message_id: Deleted message ID
        """
        if conversation_id in self.active_connections:
            deletion_data = {
                "type": "message_deleted",
                "message_id": message_id
            }
            
            disconnected = []
            for websocket in self.active_connections[conversation_id]:
                try:
                    await websocket.send_json(deletion_data)
                except:
                    disconnected.append(websocket)
            
            for ws in disconnected:
                await self.disconnect_websocket(ws, conversation_id)